package com.css.dynamic.filter;

import com.alibaba.nacos.client.naming.utils.RandomUtils;
import com.css.dynamic.rule.GrayLoadBalancer;
import org.apache.commons.lang3.StringUtils;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
import org.springframework.cloud.client.loadbalancer.reactive.DefaultResponse;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.List;
import java.util.Map;

/**
 * @author 郭辉
 * @className TODO
 * @description TODO
 * @date 2022/8/20 8:40
 * @company 海康威视
 * @since 1.0.0
 */
@Component
public class MyGrayLoadBalancer extends ReactiveLoadBalancerClientFilter {
    private static final int LOAD_BALANCER_CLIENT_FILTER_ORDER = 10150;
    private LoadBalancerProperties properties;

    private GrayLoadBalancer grayLoadBalancer;

    private DiscoveryClient discoveryClient;

    public MyGrayLoadBalancer(LoadBalancerProperties properties, GrayLoadBalancer grayLoadBalancer,DiscoveryClient discoveryClient) {
        super(null, properties);
        this.properties = properties;
        this.grayLoadBalancer = grayLoadBalancer;
        this.discoveryClient = discoveryClient;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR);
        if (!"lb".equalsIgnoreCase(schemePrefix) && !url.getScheme().equalsIgnoreCase("lb")){
            return chain.filter(exchange);
        }
        ServiceInstance serviceInstance = getServiceInstance(exchange, url);

        return Mono.just(new DefaultResponse(serviceInstance)).doOnNext((response)->{
            if (!response.hasServer()){
                throw NotFoundException.create(properties.isUse404(),
                        "Unable to find instance for " + url.getHost());
            }

            URI uri = exchange.getRequest().getURI();

            // if the `lb:<scheme>` mechanism was used, use `<scheme>` as the default,
            // if the loadbalancer doesn't provide one.
            String overrideScheme = null;
            if (schemePrefix != null) {
                overrideScheme = url.getScheme();
            }

            DelegatingServiceInstance delegatingServiceInstance = new DelegatingServiceInstance(
                    response.getServer(), overrideScheme);

            URI requestUrl = LoadBalancerUriTools.reconstructURI(delegatingServiceInstance, uri);
            exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, requestUrl);
        }).then(chain.filter(exchange));
    }

    private ServiceInstance getServiceInstance(ServerWebExchange exchange, URI url) {
        List<ServiceInstance> instances = discoveryClient.getInstances(url.getHost());
        String version = exchange.getRequest().getHeaders().getFirst("VERSION");
        if (StringUtils.isEmpty(version)){//随机取一个实例
            return instances.get(RandomUtils.nextInt(instances.size()));
        }
        ServiceInstance serviceInstance = null;
        for (ServiceInstance service : instances) {
            Map<String, String> metadata = service.getMetadata();
            if (metadata != null && metadata.get("VERSION") .equalsIgnoreCase(version)){
                serviceInstance = service;
            }
        }
        return serviceInstance;
    }

    @Override
    public int getOrder() {
        return LOAD_BALANCER_CLIENT_FILTER_ORDER;
    }

}
